package practiceProjects.AESandDatabase;

/**
 * @Project ：learn_java
 * @File ：AesExample.java
 * @IDE ：IntelliJ IDEA
 * @Author ：自由如风间
 * @Date ：2025/4/30 17:15
 * @Description ：aes加密解密示例代码
 * 最开始的代码框架是从51CTO博客找的，
 * 改了好几天，终于达到了要求。
 *
 * @BUG: 但我代码都写完了，才意识到一个问题，
 * 既然我sk和iv和密码都存在数据库中，那么如果别人看出我是aes加密的话，
 * 那不就可以直接破译了吗。。。这和明文有什么区别。。。
 * 纯当练习了。
 *
 * @FUNCTION：
 *      用AES算法存储密码项，每次运行生成不同的密钥和IV；
 *      从数据库中读取已加密的密码文本和密钥和IV解密。
 *
 * <p>
 * ECB模式:
 * 不需要初始化向量，但相同的明文块经过加密后会得到相同的密文块；
 * 因此只适合加密重复率很低的文本。
 * CBC模式:
 * 则不会有这个问题，每个明文块在加密前都会与前一个密文块进行异或运算。
 * 第一个明文块通常与一个初始化向量（IV）进行异或运算。
 * 随机性和安全性更高，性能稍弱。
 * 适合加密文本和图像等。
 */


import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.SecureRandom;
import java.sql.*;
import java.util.*;

/**
 * 密码实体类，插入数据库时用
 */
class PasswordEntity {
    String site;            // 必填
    String userName;        // 选填
    String userAccount;     // 选填
    String userPassword;    // 必填
    String postscript;      // 选填
    String secretKey;       // 必填
    byte[] iv;              // 必填
}

/**
 * 解密时用的实体类
 */
class DecryptEntity {
    String userPassword;
    String encodedKey;  // 数据库中已加密的密钥字符串
    IvParameterSpec iv;
}

public class AesExample {
    private static final String ALGORITHM = "AES";
    // 指定具体模式，加密算法/模式/填充方法
    private static final String TRANSFORMATION = "AES/CBC/PKCS5Padding";

    public static void main(String[] args) throws Exception {
        // 测试实体
        PasswordEntity pwe1 = new PasswordEntity();

        // testEncrypt(pwe1);

        testDecrypt();
    }

    public static void testEncrypt(PasswordEntity pwe) throws Exception {
        // 生成密钥
        SecretKey secretKey = generateKey();
        // 生成初始化向量
        byte[] iv = generateIv();

        // secretKey转成字符串
        String encodedKey = Base64.getEncoder().encodeToString(secretKey.getEncoded());

        /* 终端输入，直接换行则存入null */
        Scanner scanner = new Scanner(System.in);
        /*
        System.out.println("请录入信息，没有请直接换行");
        System.out.print("输入密码作用于(必填)：");
        String site = scanner.nextLine().trim();
        System.out.print("输入用户名：");
        String userName = scanner.nextLine().trim();
        System.out.print("输入账号(注册时用的手机号或邮箱)：");
        String userAccount = scanner.nextLine().trim();
        System.out.print("输入密码：");
        String userPassword = scanner.nextLine().trim();
        System.out.print("输入备注：");
        String ps = scanner.nextLine().trim();
        */

        // 直接硬性填入数据
        pwe.site = "好用便签";
        pwe.userAccount = "15023451746";
        // 原始密码文本
        pwe.userPassword = "Q@>P.44K.,rXC[EB\"r>l";
        // 存储secretKey的字符串
        pwe.secretKey = encodedKey;
        pwe.iv = iv;

        pwe.userPassword = encrypt(pwe.userPassword, secretKey, iv);
    }

    /**
     * 测试解密方法
     */
    public static void testDecrypt() {
        try (Connection conn = connect()) {
            ArrayList<DecryptEntity> secretKeyAndIVEntities = retrieveKeyAndIV(conn, 1);
            for (DecryptEntity entity : secretKeyAndIVEntities) {
                // decode the base64 encoded string
                byte[] decodedKey = Base64.getDecoder().decode(entity.encodedKey);
                // rebuild key using SecretKeySpec
                SecretKey originalKey = new SecretKeySpec(decodedKey, 0, decodedKey.length, ALGORITHM);

                String decrypted = decrypt(entity.userPassword, originalKey, entity.iv.getIV());
                System.out.println("解密结果: " + decrypted);
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public static SecretKey generateKey() throws Exception {
        KeyGenerator keyGen = KeyGenerator.getInstance(ALGORITHM);
        keyGen.init(128); // 128位密钥
        return keyGen.generateKey();
    }

    public static byte[] generateIv() {
        byte[] iv = new byte[16]; // AES块大小为16字节
        SecureRandom random = new SecureRandom();
        random.nextBytes(iv);
        return iv;
    }

    /**
     * 加密原始密码字符串
     * @param plainText 密码字符串
     * @param secretKey 密钥
     * @param iv        初始化向量
     * @return 加密处理后的base64字符串
     * @throws Exception
     */
    public static String encrypt(String plainText, SecretKey secretKey, byte[] iv) throws Exception {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        IvParameterSpec ivParams = new IvParameterSpec(iv);
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivParams);
        byte[] encryptedBytes = cipher.doFinal(plainText.getBytes());
        return Base64.getEncoder().encodeToString(encryptedBytes);
    }

    /**
     * 解密加密的密码字符串
     * @param encryptedText 加密的文本
     * @param secretKey     密钥
     * @param iv            初始化向量
     * @return 解密处理后的base64字符串
     * @throws Exception
     */
    public static String decrypt(String encryptedText, SecretKey secretKey, byte[] iv) throws Exception {
        try {
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            IvParameterSpec ivParams = new IvParameterSpec(iv);
            cipher.init(Cipher.DECRYPT_MODE, secretKey, ivParams);
            byte[] decryptedBytes = cipher.doFinal(Base64.getDecoder().decode(encryptedText));
            return new String(decryptedBytes);
        } catch (Exception e) {
            System.err.println("解密过程中出现错误: " + e.getMessage());
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 连接数据库
     * @return
     */
    private static Connection connect() {
        // 暂定的sqlite数据库地址
        String url = "jdbc:sqlite:C:/Users/zhq11/Desktop/codebook.db";
        Connection conn = null;
        try {
            conn = DriverManager.getConnection(url);
        } catch (SQLException e) {
            System.out.println(e.getMessage());
        }
        return conn;
    }

    public static void insert(PasswordEntity entity) throws SQLException {
        List<String> columns = new ArrayList<>();
        List<Object> values = new ArrayList<>();

        // 必填字段
        columns.add("site");
        values.add(entity.site);

        // 选填字段处理
        if (entity.userName != null) {
            columns.add("userName");
            values.add(entity.userName);
        }

        if (entity.userAccount != null) {
            columns.add("userAccount");
            values.add(entity.userAccount);
        }

        columns.add("userPassword");
        values.add(entity.userPassword);

        columns.add("secretKey");
        values.add(entity.secretKey);

        columns.add("iv");
        values.add(entity.iv);

        if (entity.postscript != null) {
            columns.add("postscript");
            values.add(entity.postscript);
        }

        // 构建动态sql
        String sql = "INSERT INTO pswd(" +
                String.join(", ", columns) +
                ") values (" +
                String.join(", ", Collections.nCopies(columns.size(), "?")) +
                ")";

        try (Connection conn = connect();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {
            for (int i = 0; i < values.size(); i++) {
                Object value = values.get(i);
                if (value instanceof String) {
                    pstmt.setString(i + 1, (String) value);
                } else if (value instanceof byte[]) {
                    pstmt.setBytes(i + 1, (byte[]) value);
                } else if (value instanceof SecretKey) {
                    pstmt.setBytes(i + 1, ((SecretKey) value).getEncoded());
                } else {
                    pstmt.setObject(i + 1, value);
                }
            }
            pstmt.executeUpdate();
        }
    }

    /**
     * 解密sk和iv
     * @param conn
     * @param recordId
     * @return  所有加密后的密码、密钥、iv的列表
     * @throws SQLException
     */
    public static ArrayList<DecryptEntity> retrieveKeyAndIV(Connection conn, int recordId) throws SQLException {
        // 这里的sql有多种写法，可以根据userAccount、userName、id等多种条件查询
        String sql = "SELECT userPassword, secretKey, iv FROM pswd WHERE id = ?";

        ArrayList<DecryptEntity> res = new ArrayList<>();

        try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
            pstmt.setInt(1, recordId);  // 向sql语句中填入id
            try (ResultSet rs = pstmt.executeQuery()) {
                if (rs.next()) {
                    String userPassword = rs.getString("userPassword");
                    // 重建密钥和 IV
                    String encodedKey = rs.getString("secretKey");
                    byte[] ivBytes = rs.getBytes("iv");

                    IvParameterSpec iv = new IvParameterSpec(ivBytes);

                    DecryptEntity t = new DecryptEntity();
                    t.userPassword = userPassword;
                    t.encodedKey = encodedKey;
                    t.iv = iv;
                    res.add(t);
                }
            }
        }

        return res;
    }
}
